package edu.vt.cbil.core;

import edu.vt.cbil.dataset.*;
import edu.vt.cbil.util.BioToolbox;
import edu.vt.cbil.util.ConfigObject;
import edu.vt.cbil.util.Histogram;

import java.util.*;
import java.util.logging.Logger;

/**
 * @author Justin
 * @version 1.0
 * @created 12-11??-2014 21:39:46
 */
public class Preprocessor {
    private Logger m_log;
    private BioCollection p_Collection;
    private String curSampleId;
    private BioSample p_sample;
    private ConfigObject p_config;

    private Map<Long, ProbeSetRecord> NormalData;
    private Map<Long, ProbeSetRecord> TumorData;
    private double IntensityOutlierThresh;
    private double attenuationM;

    public Preprocessor(BioCollection pc, ConfigObject pcf) {
        m_log = Logger.getLogger("edu.vt.cbil.bacom2");

        p_Collection = pc;
        p_config = pcf;

        NormalData = new TreeMap<Long, ProbeSetRecord>();
        TumorData = new TreeMap<Long, ProbeSetRecord>();
    }

    public void setCollection(BioCollection p_Collection) {
        this.p_Collection = p_Collection;
    }

    public void setCurrentSample(String sampleId) {
        curSampleId = sampleId;
        p_sample = p_Collection.getSample(sampleId);
    }

    private void clearData() {
        NormalData.clear();
        TumorData.clear();
    }

    public void preprocessSamples(String sampleId, boolean doAttenuationCorrection) {
        clearData();
        setCurrentSample(sampleId);

        for (Short i : p_Collection.getChrIdArray()) {
            for (ProbeSet ps : p_sample.getAllProbeSets(i).values()) {
                NormalData.put(ps.Loci, ps.data.Normal);
                TumorData.put(ps.Loci, ps.data.Tumor);
            }
        }

        BioToolbox.genotypeCalling(NormalData, 0.22, 0.34);
        BioToolbox.genotypeCalling(TumorData, 0.15, 0.38);

        adjustIntensitySignal();
        //考虑CN部分无法进行这部分处理，所以暂时放到后面处理
//        correctCrossTalk(NormalData);
//        correctCrossTalk(TumorData);

        IntensityOutlierThresh = calcIntensityOutlierThresh(NormalData, TumorData);
        calculteAlleleSignals();

        normalize();

        if (doAttenuationCorrection) {
            attenuationM = calculateAttenuationM();
            p_sample.attenuationM = attenuationM;

            correctAlleleAttenuation(attenuationM);
        }


        calculateCopyNumber();

        removeOutlier();

        calculateMeanIntensity();
    }

    private void normalize() {
        double sum = 0;
        int count = 0;
        for (Short index : p_Collection.ChrIdInterested) {
            for (ProbeSet p : p_sample.getAllProbeSets(index).values()) {
                if (p.data.Is_Valid) {
                    double temp;
                    if (p.Is_SNPProbe) {
                        temp = p.data.alleleA + p.data.alleleB;
                    } else {
                        temp = p.data.alleleCN;
                    }
                    if (temp < 10) {
                        sum += temp;
                        count++;
                    }
                }
            }
        }
        double avg = sum / count;
        m_log.info(String.format("allele avg = %.5f", avg));

        for (Short index : p_Collection.ChrIdInterested) {
            for (ProbeSet p : p_sample.getAllProbeSets(index).values()) {
                if (p.data.Is_Valid) {
                    if (p.Is_SNPProbe) {
                        p.data.alleleA /= avg;
                        p.data.alleleB /= avg;
                    } else {
                        p.data.alleleCN /= avg;
                    }
                }
            }
        }
    }


    private void adjustIntensitySignal() {

        //background noise offset removal
        double[] minIntensity = this.getMinimumIntensity(0.00005);

        m_log.info(String.format("minimum Intensity of Normal Sample = %.1f \t minimum Intensity of Tumor Sample = %.1f \n",
                minIntensity[0], minIntensity[1]));

        for (Chromosome chr : p_sample.getAllChromosomes()) {
            for (ProbeSet p : chr.probes.values()) {
                p.translating(minIntensity);
            }
        }
    }

    private void correctCrossTalk(Map<Long, ProbeSetRecord> inputArray) {
        List<Double> CrosstalkList = new LinkedList<Double>();
        for (ProbeSetRecord ps : inputArray.values()) {
            if (ps.isGenotypeA2B2) {
                if ((ps.mean_A == 0) && (ps.mean_B == 0)) continue;

                if (ps.mean_A > ps.mean_B) {
                    CrosstalkList.add(ps.mean_B / ps.mean_A);
                } else {
                    CrosstalkList.add(ps.mean_A / ps.mean_B);
                }
            }
        }
        double meanValue = BioToolbox.mean(CrosstalkList);
        double p1 = 1.0 / (1.0 - meanValue);
        double p2 = meanValue / (meanValue - 1);

        double minValue = 0;

        for (ProbeSetRecord ps : inputArray.values()) {
            double a = p1 * ps.mean_A + p2 * ps.mean_B;
            double b = p2 * ps.mean_A + p1 * ps.mean_B;

            ps.mean_A = a;
            ps.mean_B = b;

            minValue = Math.min(minValue, Math.min(a, b));
        }
        //归一化起点到0点
//        for (ProbeSetRecord ps : inputArray.values()) {
//            ps.mean_A -= minValue;
//            ps.mean_B -= minValue;
//        }

    }

    private void calculteAlleleSignals() {

        for (Chromosome chr : p_sample.getAllChromosomes()) {
            for (ProbeSet p : chr.probes.values()) {
                p.calculateAlleleSignals();
            }
        }
    }

    private double calculateAttenuationM() {
        Histogram h = new Histogram(1024, 0, 16);

        for (Chromosome chr : p_sample.getAllChromosomes()) {
            for (ProbeSet p : chr.probes.values()) {
                if (p.data.Is_Valid) {
                    if (p.data.Tumor.mean_A < IntensityOutlierThresh) {
                        h.addDataPoint(p.data.alleleA);
                    }
                    if (p.data.Tumor.mean_B < IntensityOutlierThresh) {
                        h.addDataPoint(p.data.alleleB);
                    }
                }
            }
        }
        double[] thresh = h.getThresh(1e-6);//0.000001

        double m = thresh[1];

        if (m > p_config.params.limitAttenuationM) {
            m_log.warning(String.format("attenuation m = %.5f", m));

            m = p_config.params.limitAttenuationM;
        }

        m_log.info(String.format("attenuation m = %.5f", m));
        return m;
    }

    private void correctAlleleAttenuation(double m) {
        for (Chromosome chromosome : p_sample.getAllChromosomes()) {
            for (ProbeSet p : chromosome.probes.values()) {
                if (p.data.Is_Valid) {
                    if (p.Is_SNPProbe) {
                        p.data.alleleA = reviseAttenuation(p.data.alleleA, m);
                        p.data.alleleB = reviseAttenuation(p.data.alleleB, m);
                    } else if (p.Is_CNProbe) {
                        p.data.alleleCN = reviseAttenuation(p.data.alleleCN, m);
                    }
                }
            }
        }
    }

    private final double reviseAttenuation(double value, double m) {
//      method 1:
//      return value / (1 - (1 - value) / m);
//      method 2:
//       return value / (1 - (value - 1) / m);
        //      method 3:
        final double n = 1;
        if (value < 0.9 * m) {
            return value / (1 - (n * value - 1) / (n * m));
        } else {
            //return 0.9*m /(1 -  0.25*(4*0.9*m - 1) / m);
            return 0.9 * n * m * m / (0.1 * n * m + 1);
        }
//      method 4:
//        if (value < 0.95*m){
//            return (m-1)*value/(m - value);
//        }else{
//            return (m-1)*19;//(m-1)*0.95*m/(m - 0.95*m);
//        }
//      method 5:
//        final double n = 1.2;
//            return value / (1 - value/(n*m-1));

    }


    public void calculateCopyNumber() {
        for (Chromosome chr : p_sample.getAllChromosomes()) {
            for (ProbeSet p : chr.probes.values()) {
                p.calculateCopyNumber();
            }
        }

    }


//    public void searchExtremalCopyNumber() {
//
//        List<Double> globalMin = new ArrayList<Double>();
//        List<Double> globalMax = new ArrayList<Double>();
//
//        for (Chromosome chromosome : p_Collection.sample) {
//            if (chromosome.probes.size() == 0) continue;
//
//            int count = 0;
//            int cid = 0;
//            double minValue = 0, maxValue = 0;
//
//            for (ProbeSet p : chromosome.probes.values()) {
//                if (p.data.Is_Valid) {
//                    //
//                    cid = p.Chr_id;
//
//                    if (count == 0) {
//
//                        minValue = p.data.CopyNumberHalf;
//                        maxValue = p.data.CopyNumberHalf;
//
//                        count = 1;
//                        continue;
//                    }
//
//                    minValue = Math.min(p.data.CopyNumberHalf, minValue);
//                    maxValue = Math.max(p.data.CopyNumberHalf, maxValue);
//                    count++;
//                }
//            }
//            MinMaxCopyNumber.get(cid)[0] = minValue;
//            MinMaxCopyNumber.get(cid)[1] = maxValue;
//            globalMin.add(minValue);
//            globalMax.add(maxValue);
//        }
//
//        // Global
//        MinMaxCopyNumber.get(0)[0] = Collections.min(globalMin);
//        MinMaxCopyNumber.get(0)[1] = Collections.max(globalMax);
//    }

    public double[] getMinimumIntensity(double thresh) {
        int width = 100;
        Histogram hn = new Histogram(width, 10, 1000);
        Histogram ht = new Histogram(width, 10, 1000);

        for (Chromosome chromosome : p_sample.getAllChromosomes()) {
            if (chromosome.probes.size() == 0) continue;
            for (ProbeSet p : chromosome.probes.values()) {
                hn.addDataPoint(p.data.Normal.mean_A);
                hn.addDataPoint(p.data.Normal.mean_B);

                ht.addDataPoint(p.data.Tumor.mean_A);
                ht.addDataPoint(p.data.Tumor.mean_B);
            }
        }

//        System.out.println(hn.toString("Intensity Normal Sample:"));
//        System.out.println(ht.toString("Intensity Tumor Sample:"));

        return new double[]{getRoughMinimum(hn, thresh), getRoughMinimum(ht, thresh)};
    }

    private double getRoughMinimum(Histogram hist, double thresh) {
        int width = hist.getBin();
        int count = (int) (hist.getCount() * thresh);
        for (int i = 1; i < width; i++) {
            count -= hist.getFreq(i);
            if (count < 0) {
                return hist.getThresh(i - 1);
            }
        }
        return 0;
    }

    private double calcIntensityOutlierThresh(Map<Long, ProbeSetRecord> array1, Map<Long, ProbeSetRecord> array2) {
        int count = 2 * (array1.size() + array2.size() + 1);

        double[] alleleIntensity = new double[count];

        for (ProbeSetRecord pr : array1.values()) {
            if (pr.mean_CN == 0) {  // exclude CN Loci
                alleleIntensity[--count] = pr.mean_A;
                alleleIntensity[--count] = pr.mean_B;
            }
        }
        for (ProbeSetRecord pr : array2.values()) {
            if (pr.mean_CN == 0) {
                alleleIntensity[--count] = pr.mean_A;
                alleleIntensity[--count] = pr.mean_B;
            }
        }
        alleleIntensity = Arrays.copyOfRange(alleleIntensity, count - 1, alleleIntensity.length);
        Arrays.sort(alleleIntensity);
        int pos = (int) (alleleIntensity.length * 0.95);
        double thresh = alleleIntensity[pos];

        if (thresh < 6000) {
            m_log.warning(String.format("Outlier Thresh of Intensity = %.5f", thresh));

            thresh = 6000;
        }

        m_log.info(String.format("Outlier Thresh of Intensity = %.5f", thresh));

        return thresh;
    }

    private Histogram calcCopyNumberHistogram() {
        Histogram hist = new Histogram(32, 0, 8);

        for (Chromosome chromosome : p_sample.getAllChromosomes()) {
            for (ProbeSet p : chromosome.probes.values()) {
                if (p.data.Is_Valid) {
                    hist.addDataPoint(p.data.CopyNumberHalf);
                }
            }
        }
        m_log.info(hist.toString("Half CopyNumber"));
        return hist;
    }

    private void removeOutlier() {

        Histogram hist = calcCopyNumberHistogram();

        double r = 0.000001;

        double[] CNThresh = hist.getThresh(r);
        m_log.info(String.format("Thresh of Histogram [ %.5f - %.5f ] \n", CNThresh[0], CNThresh[1]));
        CNThresh[0] = Math.max(-1, CNThresh[0]);
        CNThresh[1] = Math.min(5, CNThresh[1]);

        for (Chromosome chromosome : p_sample.getAllChromosomes()) {
            for (ProbeSet p : chromosome.probes.values()) {
                if (p.data.Is_Valid) {
                    if ((p.data.CopyNumberHalf < CNThresh[0])
                            || (p.data.CopyNumberHalf > CNThresh[1])) {
                        p.data.Is_Valid = false;
                        continue;
                    }
                    if (p.Is_SNPProbe) {
                        if (p.data.Tumor.mean_A > IntensityOutlierThresh) {
                            p.data.Is_Valid = false;
                            continue;
                        }
                        if (p.data.Tumor.mean_B > IntensityOutlierThresh) {
                            p.data.Is_Valid = false;
                            continue;
                        }
                        if (p.data.Normal.mean_A > IntensityOutlierThresh) {
                            p.data.Is_Valid = false;
                            continue;
                        }
                        if (p.data.Normal.mean_B > IntensityOutlierThresh) {
                            p.data.Is_Valid = false;
                            continue;
                        }
//                        if (p.data.Tumor.mean_A + p.data.Tumor.mean_B +
//                                p.data.Normal.mean_A + p.data.Normal.mean_B < 2) {
//                            p.data.Is_Valid = false;
//                            continue;
//                        }
//                        if ((p.data.alleleA > 3.5) || (p.data.alleleB > 3.5)) {
//                            p.data.Is_Valid = false;
//                            continue;
//                        }
                    } else if (p.Is_CNProbe) {
                        if ((p.data.Tumor.mean_CN > IntensityOutlierThresh) || (p.data.Tumor.mean_CN < 0)) {
                            p.data.Is_Valid = false;
                            continue;
                        }
                        if ((p.data.Normal.mean_CN > IntensityOutlierThresh) || (p.data.Normal.mean_CN < 0)) {
                            p.data.Is_Valid = false;
                            continue;
                        }

                    }
                }
            }
        }


    }

    public void calculateMeanIntensity() {

        double[] sum = new double[8];
        Arrays.fill(sum, 0);
//        sum[0] = 0;
//        sum[1] = 0;
//        sum[2] = 0;
//        sum[3] = 0;
//        sum[4] = 0;
//        sum[5] = 0;
//        sum[6] = 0;
//        sum[7] = 0;
        int count = 0;
        for (Chromosome chr : p_sample.getAllChromosomes()) {
            for (ProbeSet p : chr.probes.values()) {
                if (p.data.Is_Valid) {
                    sum[0] += p.data.Tumor.mean_A;
                    sum[1] += p.data.Tumor.mean_B;
                    sum[2] += p.data.Normal.mean_A;
                    sum[3] += p.data.Normal.mean_B;
                    sum[4] += p.data.alleleA;
                    sum[5] += p.data.alleleB;
                    sum[6] += p.data.CopyNumberHalf;
                    count++;
                }

            }
        }
        m_log.info(String.format("Loci count = %d \t Avg -> Tumor A = %.5f \t Tumor B = %.5f \t Normal A = %.5f \t Normal B = %.5f \t Allele A = %.5f \t Allele B = %.5f \t CN = %.5f",
                count, sum[0] / count, sum[1] / count, sum[2] / count, sum[3] / count, sum[4] / count, sum[5] / count, sum[6] / count));
    }


}